#!/usr/bin/ruby

require "bio"
require "ruby-debug"
Debugger.start

class MatchMatrices

	attr_accessor :matrices

	def initialize(file)
		@matrix_file = File.open(file)

		file_dat = @matrix_file.read.split("//")
		file_dat.delete_at(file_dat.length - 1)
		@matrices = file_dat.map {|mat| MatchMatrix.new(mat)}

	end


	def quick_parse
		temp_mat = []
		@matrices.each do |mat|
			if(mat.attrib["NA"] =~ /E2F/)
				temp_mat << mat
			else
				temp_mat << mat if(rand(100) == 0)
			end
		end
		@matrices = temp_mat
	end

	def match_seq(seq)
		seq = Bio::Sequence::NA.new(seq) if(seq.is_a?(String))
		ret_hash = {}

		@matrices.map do |mat| 
			ret_hash[mat.attrib["ID"]] = mat.match(seq)
		end

		return ret_hash
	end

	def to_s
		@matrices.join("")
	end


	class MatchMatrix

		attr_reader :weights, :attrib

		def initialize(mat_s)
			mat_dat = mat_s.split("WEIGHTS")

			@attrib = {}
			mat_dat[0].strip.split("\n").each do |attrib|
				attr_sp = attrib.split(" ")
				@attrib[attr_sp[0]] = attr_sp[1]
			end

			@weights = mat_dat[1].strip.split("\n").map do |weight_line|
				weight_array = weight_line.split(" ")
				weight_array.delete_at(0)

				weight_hash = {:N => 0}
				weight_array.each do |w|
					w_p = w.split(":")
					# Create the hash
					weight_hash[w_p[0].to_sym] = w_p[1].to_f
				end

				weight_hash
			end

		end

		def seq_to_sym(seq)
			seq.to_s.split("").map {|base| base.upcase.to_sym}
		end

		def match(seq)
			# debugger
			puts "Matching " + @attrib["ID"]
			forward_sym = seq_to_sym(seq)
			back_sym = seq_to_sym(seq.complement)

			f_res = match_sym(forward_sym).max
			b_res = match_sym(back_sym).max
			if(f_res > b_res)
				res = f_res
			else
				res = b_res
			end
			return res
		end

		def match_sym(sym_seq)
			frames = sym_seq.size - @weights.size
			return [0] if(frames < 0)
			mat_max 	= @attrib["MAXIMAL"].to_f 

			max 		= 0
			results 	= []
			0.upto(frames) do |seq_loc|
				frame_total = 0
				@weights.each_with_index do |weight, mat_loc|
					#debugger
					frame_total += weight[sym_seq[seq_loc + mat_loc]]
				end
				max = frame_total if(frame_total > max)
				results << frame_total / mat_max
			end
				
			return results
		end

		def to_s
			@weights.map {|w| w.values.join("\t")}.join("\n")
		end

	end

end

def analyze_file(mm, file)

	fastafile = Bio::FlatFile.auto(file)
	entry_count = fastafile.entries.size
	fastafile.rewind
	count = 0

	ret_hash = {}
	fastafile.each_entry do |entry|
		count += 1
		percent = ((count.to_f / entry_count.to_f) * 100).to_i
		puts "\n#{percent}%: ANALYZING " + entry.entry_id.to_s
		ret_hash[entry.entry_id] = mm.match_seq(entry.data.gsub("\n", "").slice(0, 2000))
	end


	ret_hash
end

def write_hash(hash, file_out)

	f = File.open(file_out, 'w')
	int_keys = hash.values[0].keys

	f.write( "Gene\t" + int_keys.join("\t") )

	hash.each_pair do |hk, hv|
		s = hk
		s = "" unless(s)
		int_keys.each {|ik| s += "\t" + hv[ik].to_s}
		f.write( "\n" + s )
	end

	f.close
end

case(ARGV.length)
when 0
when 1
	puts MatchMatrices.new(ARGV[0]).to_s

when 2..4
	mm = MatchMatrices.new(ARGV[0])
	if(ARGV.length < 4)
		write_hash(analyze_file(mm, ARGV[1]), ((ARGV.length > 2) ? ARGV[2] : "out.txt"))
	end
else
	puts "Invalid Arguments:\n"
end
